The official implementation can be found at https://github.com/huangdaolang/amortized-decision-aware-bed